import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.autograd import Variable


# 读取数据并对数据做预处理,画出这些数据在图上的点
with open('./Data/data.txt', 'r') as f:
    data_list = f.readlines()
    # print(data_list)
    # print()
    # data_list = [i.split('\n')[0] for i in data_list]
    temp_lst = []
    for i in data_list:
        temp_lst.append(i.split('\n')[0])
    data_list = temp_lst
    # print(data_list)
    # print()
    temp_lst = []
    # data_list = [i.split(',') for i in data_list]
    for i in data_list:
        temp_lst.append(i.split(','))
    data_list = temp_lst
    # print(data_list)
    # print()
    data = [(float(i[0]), float(i[1]), float(i[2])) for i in data_list]


# 将标志位为0的数据放在x0中，标志位为1的数据放在x1中
x0 = list(filter(lambda x: x[-1] == 0.0, data))
x1 = list(filter(lambda x: x[-1] == 1.0, data))

plot_x0_x = [i[0] for i in x0]
plot_x0_y = [i[1] for i in x0]
plot_x1_x = [i[0] for i in x1]
plot_x1_y = [i[1] for i in x1]

plt.plot(plot_x0_x, plot_x0_y, 'ro')
plt.plot(plot_x1_x, plot_x1_y, 'bo')


# plt.show()


# 定义设计模型样式
# 使用了一个全连接层,输入为2维数据，输入为1一维数据(同时也可以理解为前一层神经网络为两个节点向这一层的一个节点输入数据并输出)
# 使用了一个Sigmoid激活函数
class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.lr = nn.Linear(2, 1)
        self.sm = nn.Sigmoid()

    def forward(self, x):
        x = self.lr(x)
        x = self.sm(x)
        return x


# 将模型实例化
model = LogisticRegression()
criterion = nn.BCELoss()  # 二分类交叉熵函数
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)  # 加了惯性(推动力)的梯度下降算法

# 处理得出训练集的数据x_data以及y_data
np_data = np.array(data, dtype='float32')  # 转换成 numpy array
x_data = torch.from_numpy(np_data[:, 0:2])  # 转换成 Tensor, 大小是 [100, 2]
y_data = torch.from_numpy(np_data[:, -1]).unsqueeze(1)  # 转换成 Tensor，大小是 [100, 1]

# 开始训练模型
epochs = 20000
for epoch in range(epochs):
    x = Variable(x_data)
    y = Variable(y_data)  # 将训练集数据以及label数据转换为Variable类型数据

    # 开始编辑正向传播部分的逻辑
    out = model(x)  # 将训练集数据放入模型中获得预测到的数据
    loss = criterion(out, y)  # 用模型预测出的值与label数据进行交叉熵计算，计算出对应的损失值
    print_loss = loss.data.numpy()  # print_loss变量存放一下损失值，便于后面直接输出打印
    mask = out.ge(0.5).float()  # 将out中所有的元素与0.5比较，大于0.5时为1，小于0.5时为0(二分法问题，大于0.5判定为1小于0.5判定为0)
    correct = (mask == y).sum()  # 对mask与label中的值相同的数据数统计一下，存放到correct变量中，表示的是预测正确的数量
    acc = correct.numpy() / x.size(0)  # 正确数 / 所有数据数量 = 正确率

    # 开始编辑反向传播部分的逻辑
    optimizer.zero_grad()  # 梯度清零，防止梯度累加
    loss.backward()  # 对损失函数求导并进行反向传播求出的梯度值
    optimizer.step()  # 修正优化器中的学习率

    if (epoch + 1) % 1000 == 0:
        print('*' * 10)
        print('epoch:{}  loss:{}  acc:{}'.format((epoch + 1), print_loss, acc))

# 将训练好的二分线显示在点状图中查看效果
w0, w1 = model.lr.weight[0]
w0 = w0.data.numpy()
w1 = w1.data.numpy()
b = model.lr.bias.data.numpy()
plot_x = np.arange(30, 100, 0.1)
plot_y = (-w0 * plot_x - b) / w1
plt.plot(plot_x, plot_y)
plt.show()
